#!/usr/bin/env python
# -*- coding:utf-8 -*-
from conf.setting import DICT_PATH
import os
import traceback
try:
    from jwt import decode,InvalidTokenError,DecodeError,get_unverified_header
    from tqdm import tqdm
except ImportError:
    print("请安装扩展模块,pyjwt、tqdm!")

def get_files(fDir,suf):
    '''
    @des
    获取指定目录下指定文件格式的文件名
    suf为空，则获取完整文件名
    '''
    result = []
    for _,_,files in os.walk(fDir):
        for f in files:
            name,suffix = os.path.splitext(f)
            if suf:
                if suffix == suf:
                    result.append(name)
            else:
                result.append(f)
    return result

def is_jwt(content:str):
    plist = content.split(".")
    if len(plist) != 3:
        return False
    return True

def judge_sign_type(content:str):
    header = get_unverified_header(content)
    if header['alg'] in ["HS256","HS512","HS384"]:
        return True
    return False
def is_alg(content:str):
    return get_unverified_header(content)['alg']

def find_jwt_secret(content:str,dictFile:str):
    htype = is_alg(content)
    with open(dictFile,encoding="utf-8") as fp:
        for f in tqdm(fp,desc="破解进度"):
            f = f.strip()
            try:
                decode(content,f,algorithms=htype)
                return f
            except DecodeError:
                pass
            except InvalidTokenError:
                return f
    return False

def execute(args):
    try:
        print("工具收集的字典如下:")
        dfiles = get_files(DICT_PATH,"")
        cdict = ''
        if dfiles:
            for d in dfiles:
                print(d)
            cdict = input("请输入要使用的字典: ")
            if cdict in dfiles:
                cdict = DICT_PATH + os.sep + cdict
                jwt = input("请输入jwt的值: ")
                if is_jwt(jwt):
                    result = find_jwt_secret(jwt,cdict)
                    if result:
                        print("找到签名密钥: " + result)
                    else:
                        print("未找到签名密钥，建议更换字典!")
                else:
                    print("jwt格式错误!")
            else:
                print("字典加载错误!")
        else:
            print("未获取到字典!")
    except Exception as e:
        traceback.print_exc()

